#!/usr/bin/python
#coding:utf-8

# Copyright 2012 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0] == 'extract_edgels-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from numpy import dot

from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors

from quaternion import Quat, random_quaternion

import simplejson

import Image

set_printoptions(precision=3)


def create_pic(job_params, framenum):
    filename = job_params['root_directory']+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], double(job_params["gaussian_smoothing_factor"]))


    ## Creates picture object
    if not 'projection_model' in job_params:
        model = 'perspective'
    else:
        model = job_params['projection_model']

    if model == 'perspective':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance_guess']
        p_point = array(job_params['principal_point_guess'])
        # focal_distance = job_params['focal_distance']
        # p_point = array(job_params['principal_point'])

        pic = PicturePinhole(imr, focal_distance, p_point[0], p_point[1])

    elif model == 'equirectangular':
        pic = PictureEqrec(imr)

    return pic

if __name__ == '__main__':

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0] == 'extract_edgels-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load images and initialize pic objects

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    framenum = int(sys.argv[2])

    pic = create_pic(job_params, framenum)

    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold'] # ** -1
    ## Run the grid sweep, and generate the edgels list.
    ## Edgel extractor parameters
    pic.extract_edgels(gspc, glim)

    pic.calculate_edgel_normals()
    pic.generate_random_edgels()
    pic.calculate_random_normals()

    ## Run optimization
    sig = job_params['angular_threshold']

    #xini=[random_quaternion() for x in range(20)] ## Lots of random initial points
    #xini = array(job_params['initial_guesses'])   ## Read initial points form job file
    #xini = array([[1.,0,0,0]]) ## Just "no rotation"
    xini = []

    xopt, sig_hat = pic.find_orientation(xini, sig, method=0)
    #xopt, val = pic.find_orientation(xini, sig, passes=1)
    ##
    #################################################################
    pic.orientation = xopt.canonical()
    print 'Ans:', pic.orientation

    sig=sig_hat#.1
    pic.set_labels(sig)
    vps = pic.find_vanishing_points(sig)
    

    if PlotStuff:
        ## Plot estimated vanishing point directions
        figure(1, figsize=(12,9))

        title('Possible edge directions at multiple points')
        imshow(pic.frame/260., interpolation='nearest')
        pic.plot_vdirs(gca(), 50, pic.orientation.rot().T )
        axis('equal')
        savefig(job_params['root_directory']+'/movie/vps-%04d.png'%framenum, dpi=150)

        figure(2, figsize=(12,9))

        title('Extracted edgels')
        imshow(pic.frame/260., interpolation='nearest')

        pic.plot_edgels(gca(),5)

        axis('equal')
        savefig(job_params['root_directory']+'/movie/vps-%04d.png'%framenum, dpi=150)

        # figure(5)
        # plot_aux.plot_spheres(xopt1.canonical().rot().T, pic.normals )
        figure(6)
        plot_aux.plot_spheres(xopt.canonical().rot().T, pic.normals )
        # figure(27)
        # pp = mgrid[0.:rezl1.shape[0]]/rezl1.shape[0]
        # plot(rezl1, pp)
        # plot(reztu, pp)
        # grid()




